# src/plot_data/simpleplot.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*

import os
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

# 加入拐点
# 1. 通过时间数据找到拐点，输入时间一个接口的list时间，输出拐点的进程数
# 2. 将拐点进程数红色标记
def find_totaltime_turningpoint(data):
    totaltime = data
    for i in range(1,len(totaltime)):
        if(totaltime[i] >= totaltime[i-1]):
            break
    return i


# myplot
# --figtype TotalTime Sizeperprocess ParallelEfficiency TimePercent ComponentTime
# Timepercent 才有 --CompareObj Upper TotalTime
# Timepercent ComponentTime 才有 --level

# 输入 X Y ylabel legend_name title
def plot_linechart(savedir,x,y,ylabel,legend_name,title):
    """plot a line chart"""
    length = len(x)
    x_axis = np.linspace(1, length, length,endpoint=False)#首先等距生成横轴坐标
    #TODO:这部分代码这么写，连一维图像都画不了
    if(np.array(y).ndim == 1):
        plt.plot(x_axis,y, marker = "o")
    else:
        for y_axis in y :
            plt.plot(x_axis,y_axis, marker = "o")
    plt.xticks(x_axis,x)#替换掉横轴的刻度
    plt.xticks(fontsize = 8)#控制横轴刻度字体大小
    plt.yticks(fontsize = 8)#控制纵轴刻度字体大小
    plt.xlabel(r"proc_num",fontsize = 16)#横轴标签
    plt.ylabel(f"{ylabel}",fontsize = 16)#纵轴标签
    plt.title(title,fontsize = 26)#图表题目
    #legend_name = legend_name.split()#图例输入默认为字符串，且以空格间隔
    plt.legend(legend_name,fontsize = 6)#图例
    savedir = os.path.join(savedir, title +"-linechart.png")
    plt.savefig(savedir)#存储名称，这里我认为可以多加一个输入参数mydir,这样就方便确定生成图片的位置了
    plt.close()
    # plt.show()
    return

# 输入 X Y ylabel legend_name title

def plot_percent(savedir,x,y,ylabel,legend_name,title,Total_time = [],tagturningpoint=0):
    """plot a percent stacked bar"""
    x_axis = range(len(x))
    width = 0.35#柱状图宽度
    bottom_y = np.zeros(len(x))
    Totaltime_id = -1
    if ylabel == "percent_total" :
        Totaltime_id = 0
        sums = np.array(Total_time[Totaltime_id])
    else :
        sums = np.sum(y,axis=0)#only for test，这里我假设了输入y是numpy而且还没有换算到百分比，如果y是百分比那么可以删去
    for i in y :
        z = i / sums * 100#only for test，计算百分比，如果y就是百分比数据，删去
        plt.bar(x_axis,z,width,bottom=bottom_y)#remeber to change z to i
        bottom_y = z + bottom_y#remeber to change z to i
        if Totaltime_id != -1 :
            Totaltime_id += 1
            sums = np.array(Total_time[Totaltime_id])
    plt.xticks(x_axis, x)
    plt.xticks(fontsize=8)
    plt.yticks(fontsize=8)
    plt.xlabel(r"proc_num", fontsize=16)
    plt.ylabel(f"{ylabel}(%)", fontsize=16)
    plt.title(title, fontsize=26)
    #legend_name = legend_name.split()
    plt.legend(legend_name, fontsize=6, loc="upper right")
    if tagturningpoint:
        turningpointindex = find_totaltime_turningpoint(Total_time)
        ax = plt.gca()
        rect = plt.Rectangle((x_axis[turningpointindex] - 0.35/2 - 0.2,0), 0.7, max(bottom_y) + 3, fill=False, edgecolor = 'red', linewidth=2)
        ax.add_patch(rect)
    if ylabel == "percent_par" :
        savedir = os.path.join(savedir, title + "-percentstackedbar.png")
        plt.savefig(savedir)#同理可以加一个mydir确定输出文件位置
    else :
        savedir = os.path.join(savedir, title + "-percenttotal.png")
        plt.savefig(savedir)
    plt.close()
    print(f">> save pic in {savedir}")
#    plt.show()
    return


def plot_percent_2D(savedir,x,y,ylabel,legend_name,title,Total_time):
    """plot a percent stacked bar"""
    y = y.drop(labels="np",axis=1)
    x_axis = range(len(x))
    width = 0.35#柱状图宽度
    bottom_y = np.zeros(len(x))
    y = y / Total_time
    for i in list(y.columns) :
        plt.bar(x_axis,y[i]*100,width,bottom=bottom_y)#remeber to change z to i
        bottom_y = y[i] + bottom_y#remeber to change z to i
    #     z = i / sums#only for test，计算百分比，如果y就是百分比数据，删去
    #     plt.bar(x_axis,z,width,bottom=bottom_y)#remeber to change z to i
    #     bottom_y = z + bottom_y#remeber to change z to i
    #     if Totaltime_id != -1 :
    #         Totaltime_id += 1
    #         sums = np.array(Total_time[Totaltime_id])
    plt.xticks(x_axis, x)
    plt.xticks(fontsize=8)
    plt.yticks(fontsize=8)
    plt.xlabel(r"proc_num", fontsize=16)
    plt.ylabel(f"{ylabel}(%)", fontsize=16)
    plt.title(title, fontsize=26)
    #legend_name = legend_name.split()
    plt.legend(legend_name, fontsize=6,loc="upper right")
    savedir = os.path.join(savedir, title + "-percenttotal.png")
    plt.savefig(savedir)
    plt.close()
#    plt.show()
    return

def plot_percent_line(savedir,x,y,ylabel,legend_name,title,Total_time):
    """每个算例占当前进程数的一级接口 (Setup, Solve) 百分比折线图"""
    #TODO:这里都不需要输入X和legend_name，直接从数据中拿出来
    y = y.drop(labels="np",axis=1)
    length = len(x)
    x_axis = np.linspace(1, length, length,endpoint=False)#首先等距生成横轴坐标
    y = y / Total_time
    for i in list(y.columns) :
        plt.plot(x_axis,y[i] * 100, marker = "o")
    plt.xticks(x_axis,x)#替换掉横轴的刻度
    plt.xticks(fontsize = 8)#控制横轴刻度字体大小
    plt.yticks(fontsize = 8)#控制纵轴刻度字体大小
    plt.xlabel(r"proc_num",fontsize = 16)#横轴标签
    plt.ylabel(f"{ylabel}(%)",fontsize = 16)#纵轴标签
    plt.title(title,fontsize = 26)#图表题目
    #legend_name = legend_name.split()#图例输入默认为字符串，且以空格间隔
    plt.legend(legend_name,fontsize = 6)#图例
    savedir = os.path.join(savedir, title +f"-percent_line.png")
    plt.savefig(savedir)#存储名称，这里我认为可以多加一个输入参数mydir,这样就方便确定生成图片的位置了
    plt.close()
    print(f">> svae into {savedir}")
    return

def getmatrixsize():
    matrixsizedict = {}
    matrixsizedict["jxpamg_solver01"] = 2097152
    matrixsizedict["jxpamg_solver02"] = 6291456
    matrixsizedict["jxpamg_solver03"] = 83073
    matrixsizedict["jxpamg_solver04"] = 2081541
    matrixsizedict["jxpamg_lixue_matrix"] = 12075090
    matrixsizedict["jxpamg_shangfei_iter00001_job0"] = 19637808
    return matrixsizedict

# 对efficiency处理
def getplteffdata(x,y) :
    """get data use for plot efficiency line chart"""
    length = len(x)
    temp_y = []
    for i in y :
        temp_i = []
        for j in range(1,length,1) :
            temp_i.append((i[0]/i[j])/(x[j]/x[0])* 100)
        temp_y.append(temp_i)
    y = temp_y 
    x.remove(x[0])            
                #print("y",j)
                #first_num = np.
    return x,y

def plot_sizeperprocess(savedir,basename,x,title) :
    matrixsizedict = getmatrixsize()
    matrixsize = np.array(matrixsizedict[basename])
    y = matrixsize / x
    y = np.round(y,decimals=2)
    length = len(y)
    x_axis = np.linspace(1, length, length,endpoint=False)#首先等距生成横轴坐标
    plt.plot(x_axis,y, marker = "o")
    for a,b in zip(x_axis,y) :
        plt.text(a,b,b,fontsize=6)
    plt.xticks(x_axis,x)#替换掉横轴的刻度
    plt.xticks(fontsize = 8)#控制横轴刻度字体大小
    plt.yticks(fontsize = 8)#控制纵轴刻度字体大小
    plt.xlabel(r"proc_num",fontsize = 10)#横轴标签
    plt.ylabel("sizeperprocess",fontsize = 10)#纵轴标签
    plt.title(title,fontsize = 16)#图表题目
    # #legend_name = legend_name.split()#图例输入默认为字符串，且以空格间隔
    # plt.legend(legend_name,fontsize = 6)#图例
    savedir = os.path.join(savedir, "sizeperprocess-linechart.png")
    plt.savefig(savedir)#存储名称，这里我认为可以多加一个输入参数mydir,这样就方便确定生成图片的位置了
    plt.close()
    return

# 处理casename名字过长
def casename2number(casename):
    basename = os.path.basename(casename)
    basename = basename.split("-")[0]
    if(basename == "jxpamg_solver01"):
        return "C1"
    elif(basename == "jxpamg_solver02"):
        return "C2"
    elif(basename == "jxpamg_solver03"):
        return "C3"
    elif(basename == "jxpamg_solver04"):
        return "C4"
    elif(basename == "jxpamg_lixue_matrix"):
        return "C5"
    elif(basename == "jxpamg_shangfei_iter00001_job0"):
        return "C6"

#TODO: 导出为CSV 代码还需整理，合并成为一个接口


#TODO：导出热力图
def plot_heatmap(savedir,title,ylabel,data):
    plt.figure(dpi=120)
    sns.set_context({"figure.figsize":(10,5)})
    sns.heatmap(data=data,square=True)
    plt.title(title)
    plt.ylabel(ylabel)
    savedir = os.path.join(savedir, f"{title}.png")
    plt.savefig(savedir)
    plt.close()
# for test
# x = np.array(["a", "b", "c", "d"])
# y = np.array([[2, 2 ,3,2],[1,1,1,9]])
# ylabel = "xx xxx"
# title = "xxxxx"
# legend_name = "dd ddd"
# plot_linechart(x,y,ylabel,legend_name,title)
# plot_percent(x,y,ylabel,legend_name,title)
# export_csv(x,y,ylabel,legend_name,title)

defaul_fontfamily="Times New Roman"
plot_setting_dict = {
    "dpi":300,
    "fig_size":(8,4),
    "line_width":1,
    "markersize":2,
    "label_fontdict":{
        'family' : defaul_fontfamily, 
        'size'   : 12
    },
    "ticks_fontdict":{
        'family' : defaul_fontfamily, 
        'size'   : 7
    },
    "title_fontdict":{
        'family' : defaul_fontfamily, 
        'size'   : 9
    },
    "legend_fontdict":{
        'family' : defaul_fontfamily, 
        'size'   : 8
    }
}

#TODO:
# 画出宗毅那种图
# Test

def create_subplot(nrows, ncols, xlable, ylable):
    fig, axes = plt.subplots(nrows, ncols, sharex=True,figsize=plot_setting_dict["fig_size"])
    # 添加全局唯一坐标方法一：
    fig.add_subplot(111, frameon=False)
    plt.tick_params(labelcolor='none', which='both', top=False, bottom=False, left=False, right=False)
    plt.xlabel(f"{xlable}",fontdict=plot_setting_dict["label_fontdict"])#横轴标签
    plt.ylabel(f"{ylable}",fontdict=plot_setting_dict["label_fontdict"])#纵轴标签

    # 添加全局唯一坐标方法二：
    # fig.text(0.5, 0.04, 'np', ha='center')
    # fig.text(0.04, 0.5, 'time(s)', va='center', rotation='vertical')
    return fig, axes

def close_sublot(fig,ax,savedir,filename,plttype):
    savedir = os.path.join(savedir, filename +f"-{plttype}.png")
    handles, labels = ax[0,0].get_legend_handles_labels()
    fig.legend(handles, labels, ncol=6, loc='upper center', borderaxespad=0, frameon=False,prop=plot_setting_dict["legend_fontdict"])
    # plt.legend(ncol=4, bbox_to_anchor=(0, 2.45), loc='best', borderaxespad=0, frameon=False,prop=plot_setting_dict["legend_fontdict"])
    plt.tight_layout()
    plt.savefig(savedir,dpi=plot_setting_dict["dpi"],bbox_inches='tight', pad_inches = 0)
    plt.close()

def close_sublot_levelplot(row,col,fig,ax,savedir,filename,plttype,legend_name):
    savedir = os.path.join(savedir, filename +f"-{plttype}.png")

    #下面形成一个函数
    handles_max, _ = ax[0,0].get_legend_handles_labels()
    for i in range(row):
        for j in range(col):
            handles, _ = ax[i,j].get_legend_handles_labels()
            if(len(handles) > len(handles_max)):
                handles_max = handles
    
    fig.legend(handles, legend_name, ncol=10, loc='upper center', borderaxespad=0, frameon=False,prop=plot_setting_dict["legend_fontdict"])
    # plt.legend(ncol=4, bbox_to_anchor=(0, 2.45), loc='best', borderaxespad=0, frameon=False,prop=plot_setting_dict["legend_fontdict"])
    plt.tight_layout()
    plt.savefig(savedir,dpi=plot_setting_dict["dpi"],bbox_inches='tight', pad_inches = 0)
    plt.close()

def mypaper_subplot(ax_index,filename, data):
    # xlable_name = data.index.name
    pictitle = "({0})  {1}".format(filename.split("_")[0],filename.split("-")[-1])
    numprocess = list(data.index)
    natureorder = list(range(len(numprocess)))
    data.index = natureorder
    data.plot(ax=ax_index,marker='o',markersize=plot_setting_dict["markersize"],linewidth= plot_setting_dict["line_width"])

    ax_index.set_title(pictitle,fontdict=plot_setting_dict["title_fontdict"])

    # 让x轴按2为跨度跳过，并添加最后一个数
    # 这部分统一成接口
    plot_natureorder = natureorder[0:len(natureorder):3]
    if(plot_natureorder[-1] != natureorder[-1]):
        plot_natureorder.append(natureorder[-1])
    plot_np = numprocess[0:len(numprocess):3]
    if(plot_np[-1] != numprocess[-1]):
        plot_np.append(numprocess[-1])

    ax_index.set_xticks(plot_natureorder)
    ax_index.set_xticklabels(plot_np, fontdict=plot_setting_dict["ticks_fontdict"], minor=False)
    # ax_index.set_xlabel(f"{xlable_name}", fontdict=plot_setting_dict["label_fontdict"])#横轴标签
    
    yticks = ax_index.get_yticks().round(2)
    if(yticks[0] < 0):
        yticks = yticks[1:]
    ax_index.set_yticks(yticks)
    ax_index.set_yticklabels(yticks ,fontdict=plot_setting_dict["ticks_fontdict"])
    # ax_index.set_ylabel(f"{ylable}", fontdict=plot_setting_dict["label_fontdict"])#纵轴标签

    ax_index.tick_params(axis='both',tickdir='in',size=2)
    ax_index.legend_.remove()  # remove the individual plot legends

    # 添加网格
    ax_index.grid(linestyle = '--')

def mypaper_subplot_level(ax_index, data):
    # xlable_name = data.index.name
    numprocess = list(data.index)
    natureorder = list(range(len(numprocess)))
    data.index = natureorder
    data.plot(ax=ax_index,marker='o',markersize=plot_setting_dict["markersize"],linewidth= plot_setting_dict["line_width"])

    plot_natureorder = natureorder[0:len(natureorder):3]
    if(plot_natureorder[-1] != natureorder[-1]):
        plot_natureorder.append(natureorder[-1])
    plot_np = numprocess[0:len(numprocess):3]
    if(plot_np[-1] != numprocess[-1]):
        plot_np.append(numprocess[-1])

    ax_index.set_xticks(plot_natureorder)
    ax_index.set_xticklabels(plot_np, fontdict=plot_setting_dict["ticks_fontdict"], minor=False)
    # ax_index.set_xlabel(f"{xlable_name}", fontdict=plot_setting_dict["label_fontdict"])#横轴标签

    ax_index.tick_params(axis='both',tickdir='in',size=2)
    ax_index.legend_.remove()  # remove the individual plot legends

def mypaper_subplot_level_axset(ax_index,filename,data):

    pictitle = "({0})  {1}".format(filename.split("_")[0],filename.split("-")[-1])
    ax_index.set_title(pictitle,fontdict=plot_setting_dict["title_fontdict"])

    # 让x轴按2为跨度跳过，并添加最后一个数
    # 这部分统一成接口
    
    yticks = ax_index.get_yticks().round(2)
    if(yticks[0] < 0):
        yticks = yticks[1:]
    ax_index.set_yticks(yticks)
    ax_index.set_yticklabels(yticks ,fontdict=plot_setting_dict["ticks_fontdict"])
    # ax_index.set_ylabel(f"{ylable}", fontdict=plot_setting_dict["label_fontdict"])#纵轴标签
    # 添加网格
    ax_index.grid(linestyle = '--')

def my_paper_plot_percent(ax_index,filename, data):
    """plot a percent stacked bar"""
    # xlable_name = data.index.name
    pictitle = "({0})  {1}".format(filename.split("_")[0],filename.split("-")[-1])
    numprocess = list(data.index)
    natureorder = list(range(len(numprocess)))
    data.index = natureorder

    # 求百分比
    data = data.div(data.sum(axis=1) * (0.01), axis=0)

    data.plot(ax=ax_index, kind='bar', stacked = True)

    ax_index.set_title(pictitle,fontdict=plot_setting_dict["title_fontdict"])

    # 让x轴按2为跨度跳过，并添加最后一个数
    # 这部分统一成接口
    plot_natureorder = natureorder[0:len(natureorder):3]
    if(plot_natureorder[-1] != natureorder[-1]):
        plot_natureorder.append(natureorder[-1])
    plot_np = numprocess[0:len(numprocess):3]
    if(plot_np[-1] != numprocess[-1]):
        plot_np.append(numprocess[-1])

    ax_index.set_xticks(plot_natureorder)
    ax_index.set_xticklabels(plot_np, fontdict=plot_setting_dict["ticks_fontdict"], minor=False)
    # ax_index.set_xlabel(f"{xlable_name}", fontdict=plot_setting_dict["label_fontdict"])#横轴标签
    
    yticks = ax_index.get_yticks().round(2)
    if(yticks[0] < 0):
        yticks = yticks[1:]
    if(yticks[-1] > 100):
        yticks = yticks[:-1]
    ax_index.set_yticks(yticks)
    ax_index.set_yticklabels(yticks ,fontdict=plot_setting_dict["ticks_fontdict"])
    # ax_index.set_ylabel(f"{ylable}", fontdict=plot_setting_dict["label_fontdict"])#纵轴标签

    ax_index.tick_params(axis='both',tickdir='in',size=2)
    ax_index.legend_.remove()  # remove the individual plot legends